from agent import MapAgent
from mcp_client import MCPClientSession
import asyncio
import json
from fastapi import FastAPI, HTTPException, Request
import uvicorn
from utils.logging_config import setup_logging
from utils.timeit import timeit
from contextlib import asynccontextmanager

MAX_TOOL_CALL_TIMES = 3

logger = setup_logging()
mcp_client_session: MCPClientSession | None = None
openai_format_all_tools: list = []

@asynccontextmanager
async def lifespan(app: FastAPI):
    global mcp_client_session, openai_format_all_tools
    mcp_client_session = MCPClientSession("servers_config.json")
    await mcp_client_session.start()
    all_tools = await mcp_client_session.get_all_tools()
    openai_format_all_tools = [
        convert_custom_tool_to_openai_tool(tool) for tool in all_tools
    ]
    logger.info("MCP client started")
    yield
    await mcp_client_session.stop()
    logger.info("MCP client stopped")

app = FastAPI(
    title="MapAgent Service",
    lifespan=lifespan,
)

def convert_tool_call(tool_call) -> str:
    tool_dict = {
        "tool": tool_call.function.name,
        "arguments": json.loads(tool_call.function.arguments)
    }
    return json.dumps(tool_dict)

def convert_custom_tool_to_openai_tool(tool_instance):
    """将自定义工具转换为OpenAI工具格式"""
    return {
        "type": "function",
        "function": {
            "name": tool_instance.name,
            "description": tool_instance.description.strip(),
            "parameters": tool_instance.input_schema
        }
    }

@timeit
async def execute_map_agent(query: str | None):
    if mcp_client_session is None:
        raise RuntimeError("MCP client not initialized")

    agent = MapAgent.MapAgent()
    logger.info(f"openai_format_all_tools: {openai_format_all_tools}")
    reply = agent.invoke(query, openai_format_all_tools)
    tool_call = (
        reply.choices[0].message.tool_calls[0]
        if getattr(reply.choices[0].message, "tool_calls", None)
        else None
    )

    tool_call_times = 0
    while tool_call:
        if (tool_call_times == MAX_TOOL_CALL_TIMES):
            return False, "Maximum number of tool calls exceeded."
        converted = convert_tool_call(tool_call)
        agent.add_message("Calling a tool: " + converted)

        status, result = await mcp_client_session.process_tool_call(converted)
        logger.info(f"status = {status}, result = {result}")
        if status == False:
            return False, result

        query = "Call result: " + result
        reply = agent.invoke(query, openai_format_all_tools)
        result = reply.choices[0].message.content
        result = result.removeprefix("<think>\n\n</think>\n\n")
        if "finish" in result:
            if result.startswith("```json") and result.endswith("```"):
                # 去掉 ```json\n 和 \n```
                cleaned_result = result[len("```json"):-len("```")].strip()
            return True, cleaned_result

        tool_call = (
            reply.choices[0].message.tool_calls[0]
            if getattr(reply.choices[0].message, "tool_calls", None)
            else None
        )
        tool_call_times += 1

    final_reply = reply.choices[0].message.content
    logger.info(f"FINAL REPLY: {final_reply}")
    return False, final_reply

@app.post("/callMapAgent")
async def call_map_agent(request: Request):
    try:
        data = await request.json()
        query = data.get("query")
        status, result = await execute_map_agent(query)
        return {"status": status, "result": result}
    except Exception as e:
        logger.error(f"Error in call_map_agent: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))

def start_server():
    config = uvicorn.Config(
        app,
        host="0.0.0.0",
        port=1234,
        timeout_keep_alive=60,
        reload=False
    )
    server = uvicorn.Server(config)
    asyncio.run(server.serve())

if __name__ == "__main__":
    start_server()